package com.supporter.sysbench.work;


import com.supporter.sysbench.ConnectionPoolContext;
import com.supporter.sysbench.result.Statistics;
import com.supporter.sysbench.util.RandomData;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;

/**
 * @author xunaidong
 * @date 2021/8/8
 */
public class SQLExecutor {

    private final static Logger log = LoggerFactory.getLogger(SQLExecutor.class);
    public static int POINT_SELECTS = 10;
    public static int RANGE_SIZE = 100;
    public static int multi_cnt = 3;
    public static int BATCH_COMMIT_SIZE = 1000;
    public static int record_count = 0;
    public static boolean skip_trx = false;
    public static boolean pre_begin = true;
    public static boolean range_selects = true;

    private final static String DS_MYSQL = "mysql";
    private final static String DS_ORACLE = "oracle";

    private final Statistics statistics;
    private final String dsName;
    private final int threads;
    private ConnectionPoolContext connectionPoolContext;

    public SQLExecutor(Statistics statistics, String dsName, int threads, ConnectionPoolContext connectionPoolContext) {
        this.statistics = statistics;
        this.dsName = dsName;
        this.threads = threads;
        this.connectionPoolContext = connectionPoolContext;
    }

    /**
     * 数据准备指令
     *
     * @param index     第几张表
     * @param tableData 表数量
     */
    public void prepareData(int index, int tableData) {

        Connection conn = null;
        try {
            conn = connectionPoolContext.getConnection();
            createTable(conn, index);
            initData(conn, index, tableData);
        } catch (Exception e) {
            try {
                if (conn != null) {
                    conn.rollback();
                }
            } catch (SQLException ex) {
                log.error("初始化第{}张失败,数据回滚失败.", index, e);
            }
            log.error("初始化第{}张失败.", index, e);
        } finally {
            if (conn != null) {
                try {
                    conn.close();
                } catch (SQLException e) {
                    //ignore
                }
            }
        }
    }

    /**
     * 只读测试
     *
     * @param tableNumber 测试表的数量
     * @param tableCount  每张表的数量
     */
    public void readTest(int tableNumber, int tableCount) {
        long now = System.currentTimeMillis();
        Connection conn = null;
        try {
            conn = connectionPoolContext.getConnection();
            executeBegin(conn);
            executePointSelects(conn, tableNumber, tableCount);
            executeSimpleRanges(conn, tableNumber, tableCount);
            executeSumRanges(conn, tableNumber, tableCount);
            executeOrderRanges(conn, tableNumber, tableCount);
            executeDistinctRanges(conn, tableNumber, tableCount);
            executeCommit(conn);
            statistics.latms.addAndGet(System.currentTimeMillis() - now);
        } catch (Exception e) {
            log.error("测试只读失败.", e);
        } finally {
            if (conn != null) {
                try {
                    conn.close();
                } catch (SQLException e) {
                    //ignore
                }
            }
        }
    }

    /**
     * 读写测试
     *
     * @param tableNumber 测试表的数量
     * @param tableCount  每张表的数量
     */
    public void readAndWriteTest(int tableNumber, int tableCount) {
        long now = System.currentTimeMillis();
        Connection conn = null;
        try {
            conn = connectionPoolContext.getConnection();
            executeBegin(conn);
            executePointSelects(conn, tableNumber, tableCount);
            executeSimpleRanges(conn, tableNumber, tableCount);
            executeSumRanges(conn, tableNumber, tableCount);
            executeOrderRanges(conn, tableNumber, tableCount);
            executeDistinctRanges(conn, tableNumber, tableCount);
             executeIndexUpdates(conn, tableNumber, tableCount);
            executeNonIndexUpdates(conn, tableNumber, tableCount);
            executeDeleteInserts(conn, tableNumber, tableCount);
            executeCommit(conn);
            statistics.latms.addAndGet(System.currentTimeMillis() - now);
        } catch (Exception e) {
            log.error("测试读写失败.", e);
            try {
                conn.rollback();
            } catch (SQLException ex) {
                ex.printStackTrace();
            }
        } finally {
            if (conn != null) {
                try {
                    conn.close();
                } catch (SQLException e) {
                    //ignore
                }
            }
        }

    }

    public void createTable(Connection conn, int index) throws SQLException {
        PreparedStatement ps = null;
        try {
            StringBuffer strKey = new StringBuffer();
            if (DS_MYSQL.equals(dsName)) {
                strKey = strKey.append("PRIMARY KEY (id),KEY k_1(k))");
            }
            if (DS_ORACLE.equals(dsName)) {
                strKey.append("PRIMARY KEY (id) using index)");
            }
            ps = conn.prepareStatement("CREATE TABLE obtest" + index + "(" +
                    "id INTEGER DEFAULT '0' NOT NULL," +
                    "k INTEGER DEFAULT '0' NOT NULL," +
                    "c CHAR(120) DEFAULT '' NOT NULL," +
                    "pad CHAR(60) DEFAULT '' NOT NULL," + strKey.toString());
            ps.execute();
            ps.close();
            if (DS_ORACLE.equals(dsName)) {
                ps = conn.prepareStatement("CREATE INDEX k_1_" + index + " ON obtest" + index + "(k)");
                ps.execute();
                ps.close();
            }
            log.info("任务{} 完成表{}的创建.", Thread.currentThread().getName(), "obtest" + index);
        } catch (Exception e) {
            log.error(Thread.currentThread().getName() + " 创建表失败.", e);
            throw e;
        } finally {
            if (ps != null) {
                ps.close();
            }
        }
    }

    public void initData(Connection conn, int index, int tableData) throws SQLException {
        PreparedStatement ps = null;
        int count = tableData / BATCH_COMMIT_SIZE;
        //最后剩余多少条
        int lastCount = tableData % BATCH_COMMIT_SIZE;
        try {
            ps = conn.prepareStatement("INSERT INTO obtest" + index +
                    "(id, k, c, pad) VALUES(?,?,?,?)");
            //每BATCH_COMMIT_SIZE提交一次到数据库
            for (int i = 0; i < count + 1; i++) {
                int realCommitSize = BATCH_COMMIT_SIZE;
                if (i == count) {
                    realCommitSize = lastCount;
                }
                for (int j = 0; j < realCommitSize; j++) {
                    ps.setInt(1, i * BATCH_COMMIT_SIZE + j + 1);
                    ps.setInt(2, RandomData.getKValue(tableData));
                    ps.setString(3, RandomData.cValue());
                    ps.setString(4, RandomData.padValue());
                    ps.execute();
                }
                conn.commit();
            }
            statistics.completeTable.getAndIncrement();
            log.info("任务{} 完成表{}的数据初始化.", Thread.currentThread().getName(), "obtest" + index);
        } catch (Exception e) {
            log.error(Thread.currentThread().getName() + " 初始化数据失败.", e);
            throw e;
        } finally {
            if (ps != null) {
                try {
                    ps.close();
                } catch (SQLException e) {
                    log.error(Thread.currentThread().getName() + " 关闭SQL执行器失败.", e);
                }
            }
        }
    }

    public void executeBegin(Connection conn) {
        PreparedStatement ps = null;
        try {
            ps = conn.prepareStatement("BEGIN");
        } catch (Exception e) {
            handlerException(e, 0, 0);
        } finally {
            if (ps != null) {
                try {
                    ps.close();
                } catch (SQLException e) {
                    log.error(Thread.currentThread().getName() + " 关闭SQL执行器失败.", e);
                }
            }
        }
    }

    public void executePrepareCommit(Connection conn) {
        PreparedStatement ps = null;
        try {
            ps = conn.prepareStatement("COMMIT");
            statistics.otcnt.getAndIncrement();
            statistics.qrcnt.getAndIncrement();
        } catch (Exception e) {
            handlerException(e, 0, 0);
        } finally {
            if (ps != null) {
                try {
                    ps.close();
                } catch (SQLException e) {
                    log.error(Thread.currentThread().getName() + " 关闭SQL执行器失败.", e);
                }
            }
        }
    }

    public void executeCommit(Connection conn) {
        PreparedStatement ps = null;
        try {
            ps = conn.prepareStatement("COMMIT");
            statistics.txcnt.getAndIncrement();
        } catch (Exception e) {
            handlerException(e, 0, 0);
        } finally {
            if (ps != null) {
                try {
                    ps.close();
                } catch (SQLException e) {
                    log.error(Thread.currentThread().getName() + " 关闭SQL执行器失败.", e);
                }
            }
        }
    }

    public void executePointSelects(Connection conn, int tableNumber, int tableData) {
        for (int i = 0; i < POINT_SELECTS; i++) {
            PreparedStatement ps = null;
            ResultSet rs = null;
            int tno = RandomData.getTableNum(tableNumber, threads);
            int id = RandomData.getId(tableData);
            try {
                String sql = "select c from obtest" + tno + " where id=?";
                ps = conn.prepareStatement(sql);
                ps.setInt(1, id);
                rs = ps.executeQuery();
                rs.next();
                statistics.qrcnt.getAndIncrement();
                statistics.rdcnt.getAndIncrement();
            } catch (Exception e) {
                handlerException(e, tno, id);
            } finally {
                if (ps != null) {
                    try {
                        ps.close();
                    } catch (SQLException e) {
                        log.error(Thread.currentThread().getName() + " 关闭SQL执行器失败.", e);
                    }
                }
                if (rs != null) {
                    try {
                        rs.close();
                    } catch (SQLException e) {
                        log.error(Thread.currentThread().getName() + " 关闭SQL执行器失败.", e);
                    }
                }
            }
        }
    }

    public void executeSimpleRanges(Connection conn, int tableNumber, int tableData) {
        PreparedStatement ps = null;
        int tno = RandomData.getTableNum(tableNumber, threads);
        int id = RandomData.getId(tableData);
        try {
            ps = conn.prepareStatement("select c from obtest" + tno + " where id between ? and ?");
            ps.setInt(1, id);
            ps.setInt(2, id + RANGE_SIZE);
            ps.executeQuery();
            //读+1
            statistics.rdcnt.getAndIncrement();
            //qps+1
            statistics.qrcnt.getAndIncrement();
        } catch (Exception e) {
            handlerException(e, tno, id);
        } finally {
            if (ps != null) {
                try {
                    ps.close();
                } catch (SQLException e) {
                    log.error(Thread.currentThread().getName() + " 关闭SQL执行器失败.", e);
                }
            }
        }
    }

    public void executeSumRanges(Connection conn, int tableNumber, int tableData) {
        PreparedStatement ps = null;
        int tno = RandomData.getTableNum(tableNumber, threads);
        int id = RandomData.getId(tableData);
        try {
            ps = conn.prepareStatement("select SUM(k) from obtest" + tno + " where id between ? and ?");
            ps.setInt(1, id);
            ps.setInt(2, id + RANGE_SIZE);
            ps.executeQuery();
            //读+1
            statistics.rdcnt.getAndIncrement();
            //qps+1
            statistics.qrcnt.getAndIncrement();
        } catch (Exception e) {
            handlerException(e, tno, id);
        } finally {
            if (ps != null) {
                try {
                    ps.close();
                } catch (SQLException e) {
                    log.error(Thread.currentThread().getName() + " 关闭SQL执行器失败.", e);
                }
            }
        }
    }

    public void executeOrderRanges(Connection conn, int tableNumber, int tableData) {
        PreparedStatement ps = null;
        int tno = RandomData.getTableNum(tableNumber, threads);
        int id = RandomData.getId(tableData);
        try {
            ps = conn.prepareStatement("select c from obtest" + tno + " where id between ? and ? ORDER BY c");
            ps.setInt(1, id);
            ps.setInt(2, id + RANGE_SIZE);
            ps.executeQuery();
            //读+1
            statistics.rdcnt.getAndIncrement();
            //qps+1
            statistics.qrcnt.getAndIncrement();
        } catch (Exception e) {
            handlerException(e, tno, id);
        } finally {
            if (ps != null) {
                try {
                    ps.close();
                } catch (SQLException e) {
                    log.error(Thread.currentThread().getName() + " 关闭SQL执行器失败.", e);
                }
            }
        }
    }

    public void executeDistinctRanges(Connection conn, int tableNumber, int tableData) {
        PreparedStatement ps = null;
        int tno = RandomData.getTableNum(tableNumber, threads);
        int id = RandomData.getId(tableData);
        try {
            ps = conn.prepareStatement("select DISTINCT c from obtest" + tno + " where id between ? and ? ORDER BY c");
            ps.setInt(1, id);
            ps.setInt(2, id + RANGE_SIZE);
            ps.executeQuery();
            //读+1
            statistics.rdcnt.getAndIncrement();
            //qps+1
            statistics.qrcnt.getAndIncrement();
        } catch (Exception e) {
            handlerException(e, tno, id);
        } finally {
            if (ps != null) {
                try {
                    ps.close();
                } catch (SQLException e) {
                    log.error(Thread.currentThread().getName() + " 关闭SQL执行器失败.", e);
                }
            }
        }
    }

    public void executeIndexUpdates(Connection conn, int tableNumber, int tableData) {
        PreparedStatement ps = null;
        int tno = RandomData.getTableNum(tableNumber, threads);
        int id = RandomData.getId(tableData);
        try {
            ps = conn.prepareStatement("UPDATE obtest" + tno + " set k=k+1 where id=?");
            ps.setInt(1, id);
            ps.executeUpdate();
            statistics.wrcnt.getAndIncrement();
            statistics.qrcnt.getAndIncrement();
        } catch (Exception e) {
            handlerException(e, tno, id);
        } finally {
            if (ps != null) {
                try {
                    ps.close();
                } catch (SQLException e) {
                    log.error(Thread.currentThread().getName() + " 关闭SQL执行器失败." + tno + "::" + id, e);
                }
            }
        }
    }

    public void executeNonIndexUpdates(Connection conn, int tableNumber, int tableData) {
        PreparedStatement ps = null;
        int tno = RandomData.getTableNum(tableNumber, threads);
        int id = RandomData.getId(tableData);
        try {
            ps = conn.prepareStatement("UPDATE obtest"+tno +" set c=? where id=?");
            id = RandomData.getId(tableData);
            ps.setString(1, RandomData.cValue());
            ps.setInt(2, id);
            ps.executeUpdate();
            statistics.wrcnt.getAndIncrement();
            statistics.qrcnt.getAndIncrement();
        } catch (Exception e) {
            handlerException(e, tno, id);
        } finally {
            if (ps != null) {
                try {
                    ps.close();
                } catch (SQLException e) {
                    log.error(Thread.currentThread().getName() + " 关闭SQL执行器失败.", e);
                }
            }
        }
    }

    public void executeDeleteInserts(Connection conn, int tableNumber, int tableData) {
        PreparedStatement ps = null;
        int id = RandomData.getId(tableData);
        int k = (int) (Math.random() * tableData) + 1;
        int tno = RandomData.getTableNum(tableNumber, threads);
        String cval = RandomData.cValue();
        String pval = RandomData.padValue();
        try {
            ps = conn.prepareStatement("DELETE from obtest" + tno + " where id=?");
            ps.setInt(1, id);
            ps.executeUpdate();
            statistics.wrcnt.getAndIncrement();
            statistics.qrcnt.getAndIncrement();

        } catch (Exception e) {
            handlerException(e, tno, id);
        } finally {
            if (ps != null) {
                try {
                    ps.close();
                } catch (SQLException e) {
                    log.error(Thread.currentThread().getName() + " 关闭SQL执行器失败." + tno + "::" + id, e);
                }
            }
        }
        try {
            ps = conn.prepareStatement("INSERT INTO obtest" + tno +
                    " (id, k, c, pad) VALUES(?,?,?,?)");
            ps.setInt(1, id);
            ps.setInt(2, k);
            ps.setString(3, cval);
            ps.setString(4, pval);
            ps.executeUpdate();
            statistics.wrcnt.getAndIncrement();
            statistics.qrcnt.getAndIncrement();
        } catch (Exception e) {
            handlerException(e, tno, id);
        } finally {
            if (ps != null) {
                try {
                    ps.close();
                } catch (SQLException e) {
                    log.error(Thread.currentThread().getName() + " 关闭SQL执行器失败." + tno + "::" + id, e);
                }
            }
        }

    }
//    public void execute_bulk_insert(int tid) {
//        int tn = tid + 1;
//        try {
//            ps = sqlbench.worker[tid].conn.prepareStatement("INSERT INTO obtest" + tn +
//                    "(id, k, c, pad) VALUES(?,?,?,?)");
//            for (int i = 0; i < commit_size; i++) {
//                sqlbench.idx[tid]++;
//                ps.setInt(1, sqlbench.idx[tid]);
//                ps.setInt(2, (int) (Math.random() * (sqlbench.table_size)) + 1);
//                ps.setString(3, pubfunc.get_c_value());
//                ps.setString(4, pubfunc.get_pad_value());
//                ps.executeUpdate();
//                sqlbench.worker[tid].wrcnt++;
//                sqlbench.worker[tid].qrcnt++;
//            }
//            if (common.skip_trx) {
//                sqlbench.worker[tid].otcnt++;
//                sqlbench.worker[tid].conn.commit();
//            }
//            ps.close();
//
//        } catch (Exception e) {
//            if (sqlbench.db_debug) e.printStackTrace();
//        }
//    }
//
//    public void cmd_cleanup(Connection conn, int tid) {
//        for (int i = 0; i < sqlbench.tables; i++) {
//            if (i % sqlbench.threadcnt == tid) {
//                System.out.println("Tread-" + tid + " drop table obtest" + (i + 1));
//                drop_table(conn, i + 1);
//            }
//        }
//    }
//
//    public void drop_table(Connection conn, int n) {
//        try {
//            //PreparedStatement ps=null;
//            ps = conn.prepareStatement("DROP TABLE obtest" + n);
//            ps.execute();
//            ps.close();
//            conn.commit(); //opengauss and vastbase need commit to drop
//        } catch (Exception e) {
//            if (sqlbench.db_debug) e.printStackTrace();
//        }
//    }
//
//    public void exec_sql(Connection conn, String sql) {
//        try {
//            ps = conn.prepareStatement(sql);
//            ps.execute();
//            ps.close();
//        } catch (Exception e) {
//            if (sqlbench.db_debug) e.printStackTrace();
//        }
//    }
//

//
//
//

//

//

//

//
//    public void execute_begin(int tid) {
//        try {
//            if (pre_begin) {
//                if (!skip_trx) {
//                    ps = sqlbench.worker[tid].conn.prepareStatement("BEGIN");
//                    ps.execute();
//                    ps.close();
//                    sqlbench.worker[tid].qrcnt++;
//                    sqlbench.worker[tid].otcnt++;
//                }
//            }
//        } catch (Exception e) {
//            if (sqlbench.db_debug) e.printStackTrace();
//        }
//    }
//
//    public void execute_commit(int tid) {
//        try {
//            if (!skip_trx) {
//                ps = sqlbench.worker[tid].conn.prepareStatement("COMMIT");
//                ps.execute();
//                ps.close();
//                sqlbench.worker[tid].qrcnt++;
//                sqlbench.worker[tid].otcnt++;
//            }
//        } catch (Exception e) {
//            if (sqlbench.db_debug) e.printStackTrace();
//        }
//    }
//
//    public void execute_prepare(int tid) {
//        try {
//            if (!skip_trx) {
//                if (pre_begin) {
//                    ps = sqlbench.worker[tid].conn.prepareStatement("BEGIN");
//                    ps.execute();
//                    ps.close();
//                    sqlbench.worker[tid].qrcnt++;
//                }
//                ps = sqlbench.worker[tid].conn.prepareStatement("COMMIT");
//                ps.execute();
//                ps.close();
//                sqlbench.worker[tid].qrcnt++;
//            }
//        } catch (Exception e) {
//            if (sqlbench.db_debug) e.printStackTrace();
//        }
//    }
//

//
//    public void execute_join_selects(int tid) {
//        int id = 0;
//        String tbstr = null;
//        String tbcon = null;
//        int atn[] = new int[12];
//        for (int i = 0; i < multi_cnt; i++) atn[i] = pubfunc.get_table_num();
//        tbstr = "obtest" + atn[0] + " t" + 0;
//        tbcon = "";
//        for (int i = 1; i < multi_cnt; i++) {
//            tbstr = tbstr + ",obtest" + atn[i] + " t" + i;
//            tbcon = tbcon + "t" + (i - 1) + ".id=t" + i + ".id and ";
//        }
//        id = pubfunc.get_id();
//        try {
//            //dameng bind value issue
//            if (sqlbench.IsDameng) {
//                ps = sqlbench.worker[tid].conn.prepareStatement("select SUM(t" + (multi_cnt - 1) + ".k) from " +
//                        tbstr + " where " + tbcon + " t0.id between " + id + " and " + (id + range_size));
//            } else {
//                ps = sqlbench.worker[tid].conn.prepareStatement("select SUM(t" + (multi_cnt - 1) + ".k) from " + tbstr + " where " + tbcon + " t0.id between ? and ?");
//                ps.setInt(1, id);
//                ps.setInt(2, id + range_size);
//            }
//            ps.executeQuery();
//            ps.close();
//            sqlbench.worker[tid].rdcnt++;
//            sqlbench.worker[tid].qrcnt++;
//        } catch (Exception e) {
//            System.out.println("Conn" + tid + ":" + sqlbench.worker[tid].conn);
//            exception_handle(tid, e);
//        }
//    }
//
//    public void execute_insert(int tid) {
//        try {
//            int tn = (int) (Math.random() * (sqlbench.tables)) + 1;
//            if (sqlbench.idx[tn - 1] == 0) sqlbench.worker[tid].comm.set_tn_id(tid, tn);
//            ps = sqlbench.worker[tid].conn.prepareStatement("INSERT INTO obtest" + tn +
//                    "(id, k, c, pad) VALUES(?,?,?,?)");
//            sqlbench.idx[tn - 1]++;
//            ps.setInt(1, sqlbench.idx[tn - 1]);
//            ps.setInt(2, (int) (Math.random() * (sqlbench.table_size)) + 1);
//            ps.setString(3, pubfunc.get_c_value());
//            ps.setString(4, pubfunc.get_pad_value());
//            ps.execute();
//            sqlbench.worker[tid].qrcnt++;
//            sqlbench.worker[tid].wrcnt++;
//            ps.close();
//
//        } catch (Exception e) {
//            if (sqlbench.db_debug) e.printStackTrace();
//        }
//    }
//
//    public void set_tab_id(int tid) {
//        try {
//            for (int i = 0; i < sqlbench.tables; i++) {
//                if (i % sqlbench.threadcnt == tid) {
//                    ps = sqlbench.worker[tid].conn.prepareStatement("select max(id) from obtest" + (i + 1));
//                    rs = ps.executeQuery();
//                    rs.next();
//                    sqlbench.idx[i] = rs.getInt(1); //if null,jdbc return 0
//                    rs.close();
//                    ps.close();
//                }
//            }
//        } catch (Exception e) {
//            if (sqlbench.db_debug) e.printStackTrace();
//        }
//    }
//
//    public void set_tn_id(int tid, int tn) {
//        try {
//            ps = sqlbench.worker[tid].conn.prepareStatement("select max(id) from obtest" + tn);
//            rs = ps.executeQuery();
//            rs.next();
//            sqlbench.idx[tn - 1] = rs.getInt(1); //if null,jdbc return 0
//            rs.close();
//            ps.close();
//        } catch (Exception e) {
//            if (sqlbench.db_debug) e.printStackTrace();
//        }
//    }

    public void handlerException(Exception e, int tno, int id) {
        log.error("SQL执行异常.表obtest" + tno + "主键k:" + id, e);
        statistics.ercnt.getAndIncrement();
        throw new RuntimeException(e);
    }

}
